import os
from docx import Document
import openpyxl
import requests
from concurrent.futures import ThreadPoolExecutor, as_completed

# 1.  AI API 调用函数 
def call_ai_api(content):
    """
    调用 AI API，返回功能点分析结果。
    content: 单个段落内容
    """
    api_url = "https://open.bigmodel.cn/api/paas/v4/chat/completions"
    headers = {
        "Content-Type": "application/json",
        "Authorization": "103f067729ac42b2a159f9a66a281dd7.Q2IQv5mgWbidoacD"
    }
    
    question = f"以下这一段的软件功能点是什么？请使用短句简要回答，除了功能点以外不回复其他内容，回答内容采用,隔开: {content}"
    
    payload = {
        "model": "chatglm_6b",
        "messages": [{"role": "user", "content": question}]
    }
    
    try:
        response = requests.post(api_url, json=payload, headers=headers)
        if response.status_code == 200:
            result = response.json()
            return result.get("choices")[0].get("message").get("content")
        else:
            print(f"API请求失败，状态码: {response.status_code}")
            return None
    except Exception as e:
        print(f"调用API时出现错误: {e}")
        return None

# 2.  解析 Word 文档 
def read_and_number_headings(doc_path):
    """
    解析 Word 文档，提取标题及内容，并进行编号。
    """
    if not os.path.exists(doc_path):
        print(f"ERROR: 文件未找到 -> {doc_path}")
        return {}

    doc = Document(doc_path)
    print(f"成功打开文件 -> {doc_path}")

    heading_numbers = []  # 存储编号
    headings_dict = {}  # 存储标题和内容

    print("开始遍历标题段落...")
    last_number = None

    for para in doc.paragraphs:
        text = para.text.strip()
        style_name = para.style.name
        
        if "Heading" in style_name:
            heading_level = int(style_name.split(' ')[-1])  # 获取层级

            if heading_level > len(heading_numbers):
                for _ in range(heading_level - len(heading_numbers)):
                    heading_numbers.append(1)  
            elif heading_level < len(heading_numbers):
                heading_numbers = heading_numbers[:heading_level]
                heading_numbers[heading_level - 1] += 1
            else:
                heading_numbers[heading_level - 1] += 1

            heading_number = '.'.join(map(str, heading_numbers))

            if last_number == heading_number:
                continue

            last_number = heading_number  

            # 存储标题
            headings_dict[heading_number] = {
                "title": text,
                "content": [],
                "description": None  # 预留 AI 解析结果
            }

            last_level = heading_level  
        else:
            if heading_numbers:
                heading_number = '.'.join(map(str, heading_numbers))
                headings_dict[heading_number]["content"].append(text)
    
    return headings_dict

# 3.  并行调用 AI 处理文本内容 
def process_content_with_ai(headings_dict):
    """
    并行调用 AI API，处理标题下的内容，并存入字典。
    """
    with ThreadPoolExecutor(max_workers=10) as executor:
        future_to_heading = {}

        # 提交任务
        for heading_number, heading_info in headings_dict.items():
            if heading_info["content"]:
                content = "\n".join(heading_info["content"])  # 合并所有段落
                future = executor.submit(call_ai_api, content)
                future_to_heading[future] = heading_number

        # 收集任务结果
        for future in as_completed(future_to_heading):
            heading_number = future_to_heading[future]
            description = future.result()
            headings_dict[heading_number]["description"] = description if description else "API 请求失败"

    return headings_dict

# 4.  保存到 Excel 
def save_to_excel(headings_dict, excel_path):
    """
    将处理后的数据保存到 Excel 文件（ 按层级结构填充 ）。
    """
    wb = openpyxl.Workbook()
    ws = wb.active
    ws.title = "标题内容"

    # 设置列名
    ws.append(["一级模块", "二级模块", "三级模块", "四级模块", "五级模块", "功能点计数项名称（AI解析）"])

    row = 2  # Excel 起始行
    current_row = ["", "", "", "", "", ""]  # 用来存储各级别模块和AI解析结果

    last_level = 0  # 记录上一层级

    for heading_number, heading_info in headings_dict.items():
        title = heading_info['title']
        description = heading_info["description"]  # AI 解析结果

        # 按层级分列
        levels = heading_number.split('.')
        level = len(levels)

        #  格式化 Excel 结构 
        if level > last_level:
            #  层级增加，不换行 
            if level == 1:
                current_row[0] = title
            elif level == 2:
                current_row[1] = title
            elif level == 3:
                current_row[2] = title
            elif level == 4:
                current_row[3] = title
            elif level == 5:
                current_row[4] = title
        else:
            #  层级相同或减少，换行 
            ws.append(current_row + [description])  # 不包括编号
            row += 1  # 进入下一行
            current_row = ["", "", "", "", "", ""]  # 清空行

            #  填充新行 
            if level == 1:
                current_row[0] = title
            elif level == 2:
                current_row[1] = title
            elif level == 3:
                current_row[2] = title
            elif level == 4:
                current_row[3] = title
            elif level == 5:
                current_row[4] = title

        #  填充 AI 解析 
        current_row[5] = description

        # 更新上一层级
        last_level = level

    #  写入最后一行
    if any(current_row):
        ws.append(current_row)

    # 保存 Excel 
    wb.save(excel_path)
    print(f"数据已成功保存到 Excel -> {excel_path}")

# 5.  执行流程 
doc_path = "/Users/guhaoyu/Desktop/造价AI工作/2、客户提供的建设需求或可研文档/01.docx"
excel_path = "/Users/guhaoyu/Desktop/标题内容.xlsx"

headings_dict = read_and_number_headings(doc_path)  #  解析 Word 文档 
headings_dict = process_content_with_ai(headings_dict)  #  并行调用 AI 
save_to_excel(headings_dict, excel_path)  #  保存 Excel 
